/**
  * Created by User on 2017/4/29.
  */
package com.rep

import java.util.concurrent.TimeoutException

import akka.actor.{ActorRef, Props, ActorSystem}
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.HttpMethods._
import akka.http.scaladsl.model._
import akka.http.scaladsl.model.ws._
import akka.stream.ActorFlowMaterializer
import akka.stream.scaladsl._

import scala.concurrent.Await
import scala.concurrent.duration._

object WSRequest {

  def unapply(req: HttpRequest) : Option[HttpRequest] = {
    if (req.header[UpgradeToWebsocket].isDefined) {
      req.header[UpgradeToWebsocket] match {
        case Some(upgrade) => Some(req)
        case None => None
      }
    } else None
  }

}

/**
  * Simple websocket server using akka-http and akka-streams.
  *
  * Note that about 600 messages get queued up in the send buffer (on mac, 146988 is default socket buffer)
  */
object WSServer extends App {

  // required actorsystem and flow materializer
  // new version is flowMaterializer
  implicit val system = ActorSystem("websockets")
  implicit val fm = ActorFlowMaterializer()

  // setup the actors for the stats
  // router: will keep a list of connected actorpublisher, to inform them about new stats.
  // vmactor: will start sending messages to the router, which will pass them on to any connected routee
  val router: ActorRef = system.actorOf(Props[RouterActor], "router")
  val vmactor: ActorRef = system.actorOf(Props(classOf[VMActor], router ,2 seconds, 20 milliseconds))

  // Bind to an HTTP port and handle incoming messages.
  val binding = Http().bindAndHandleSync({
    case WSRequest(req@HttpRequest(GET, Uri.Path("/echo"), _, _, _)) => handleWith(req, Flows.echoFlow)
    case WSRequest(req@HttpRequest(GET, Uri.Path("/pushMsg"), _, _, _)) => handleWith(req, Flows.graphFlowWithExtraSource)
    case WSRequest(req@HttpRequest(GET, Uri.Path("/pushStats"), _, _, _)) => handleWith(req, Flows.graphFlowWithStats(router))
    case _: HttpRequest => HttpResponse(400, entity = "Invalid websocket request")

  }, interface = "localhost", port = 9001)

  // binding is a future, we assume it's ready within a second or timeout
  // maybe we can set it longer in real condition
  try {
    Await.result(binding, 1 second)
    println("Server online at http://localhost:9001")
  } catch {
    case exc: TimeoutException =>
      println("Server took to long to startup, shutting down")
      system.shutdown()
  }

  /**
    * Simple helper function, that connects a flow to a specific websocket upgrade request
    */
  def handleWith(req: HttpRequest, flow: Flow[Message, Message, Unit]) = req.header[UpgradeToWebsocket].get.handleMessages(flow)

}

/**
  * This object contains the flows the handle the websockets messages. Each flow is attached
  * to a websocket and gets executed whenever a message comes in from the client.
  */
object Flows {


  /**
    * Simple flow which just returns the original message back to the client
    *
    * This we show that how we get that message from client and send something back to it
    */
  def echoFlow: Flow[Message, Message, Unit] =  Flow[Message]


  /**
    * When the flow is materialized we don't really just have to respond with a single
    * message. Any message that is produced from the flow gets sent to the client. This
    * means we can also attach an additional source to the flow and use that to push
    * messages to the client.
    *
    * So this flow looks like this:
    *
    *        in ~> filter ~> merge
    *           newSource ~> merge ~> map
    * This flow filters out the incoming messages, and the merge will only see messages
    * from our new flow. All these messages get sent to the connected websocket.
    *
    * This we show that how we put message to the client
    */
  def graphFlowWithExtraSource: Flow[Message, Message, Unit] = {
    Flow() { implicit b =>
      import FlowGraph.Implicits._

      // Graph elements we'll use
      val merge = b.add(Merge[Int](2))
      val filter = b.add(Flow[Int].filter(_ => false))

      // convert to int so we can connect to merge
      val mapMsgToInt = b.add(Flow[Message].map[Int] { msg => -1 })
      val mapIntToMsg = b.add(Flow[Int].map[Message]( x => TextMessage.Strict("Hello, here is the server message")))
      val log = b.add(Flow[Int].map[Int](x => {println(x); x}))

      // source we want to use to send message to the connected websocket sink
      val rangeSource = b.add(Source(1 to 2000))

      // connect the graph
      mapMsgToInt ~> filter ~> merge // this part of the merge will never provide msgs
      rangeSource ~> log ~> merge ~> mapIntToMsg

      // expose ports
      (mapMsgToInt.inlet, mapIntToMsg.outlet)
    }
  }

  /**
    * Creates a flow which uses the provided source as additional input. This complete scenario
    * works like this:
    *  1. When the actor is created it registers itself with a router.
    *  2. the VMActor sends messages at an interval to the router.
    *  3. The router next sends the message to this source which injects it into the flow
    *
    *  This we show that how we sent message to the registered client
    */
  def graphFlowWithStats(router: ActorRef): Flow[Message, Message, Unit] = {
    Flow() { implicit b =>
      import FlowGraph.Implicits._

      // create an actor source
      // see the VMStatsPublisher in the VMStatsPublisher.scala
      val source = Source.actorPublisher[String](Props(classOf[VMStatsPublisher],router))

      // Graph elements we'll use
      val merge = b.add(Merge[String](2))
      val filter = b.add(Flow[String].filter(_ => false))

      // convert to int so we can connect to merge
      val mapMsgToString = b.add(Flow[Message].map[String] { msg => "" })
      val mapStringToMsg = b.add(Flow[String].map[Message]( x => TextMessage.Strict(x)))

      val statsSource = b.add(source)

      // connect the graph
      mapMsgToString ~> filter ~> merge // this part of the merge will never provide msgs
      statsSource ~> merge ~> mapStringToMsg

      // expose ports
      (mapMsgToString.inlet, mapStringToMsg.outlet)
    }
  }

}